net = SiameseNetwork(input_shape)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net.to(device)
criterion = torch.nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(net.parameters(),0.001, betas=(0.9,0.999))
counter = []
loss_history = []
iteration_number = 0
if __name__ == '__main__':
    for epoch in range(0, train_number_epochs):
        for i, data in enumerate(train_dataloader, 0):
            img0, img1, label = data
            img0,img1,label=img0.to(device),img1.to(device),label.to(device)
            output = net(img0, img1)
            loss_contrastive = criterion(output, label)
            optimizer.zero_grad()
            loss_contrastive.backward()
            optimizer.step()
            if i % 10 == 0:
                print("Epoch number{}\n loss {}\n".format(epoch,
                                                           loss_contrastive.item()))
                iteration_number += 10
                counter.append(iteration_number)
               loss_history.append(loss_contrastive.item())
    plt.plot(counter,loss_history)
    plt.show()
    torch.save(net.state_dict(),'weights/vgg_siamese.pt') 
